import time
import torch

from fvcore.nn import FlopCountAnalysis, parameter_count

from models.rgbtcc_mamba import get_model as rgbtcc


@torch.no_grad()
def get_flops_params(model, inputs):
    flops = FlopCountAnalysis(model, inputs)
    flops = flops.total() / 1000000000.

    params = parameter_count(model)
    params = params[""] / 1000000.

    return flops, params
    
    
if __name__ == '__main__':
    x = torch.randn(1, 3, 512, 512).cuda()
    inputs = [x, x]
    
    model = rgbtcc()
    model = model.cuda()
    model.eval()
    y = model(inputs)
    
    print('22')
    flops, params = get_flops_params(model, inputs)
    print(f"[INFO] FVcore FLOPs(G)={round(flops, 2)}, FVcore Params(M)={round(params, 2)}")
    
    for _ in range(5):
        y = model(inputs)

    n = 100
    start_time = time.time()
    for _ in range(n):
        y = model(inputs)
    end_time = time.time()
    
    fps = n / (end_time - start_time)
    print(f"[INFO] FPS={round(fps, 2)}")
